import cv2
import torch
import numpy as np
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from facelib import get_config
from facelib import add_from_folder
conf = get_config()


class FaceAnalysis:
    def __init__(self,
                 facebank_path=None,
                 device: str = None):
        """
        初始化人脸分析系统
        Args:
            facebank_path: 人脸特征库路径
            device: 计算设备(cuda/cpu)
        """

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        self.facebank_path = conf.facebank_path if facebank_path is None else Path(facebank_path)

        print("正在初始化人脸分析系统……")
        # 初始化各组件
        from facelib import FaceDetector, FaceRecognizer, EmotionDetector, AgeGenderEstimator, get_config
        self.detector = FaceDetector(device=self.device)
        self.emotion_detector = FaceDetector(face_size=(224, 224))
        self.recognizer = FaceRecognizer(conf=conf)
        self.emotion = EmotionDetector(device=self.device)
        self.age_gender = AgeGenderEstimator(device=self.device)

        # 状态跟踪
        self._last_operations = {
            'detection': False,
            'recognition': False,
            'emotion': False,
            'age_gender': False
        }

        # 添加检测结果缓存
        self._last_detection = {
            'boxes': None,
            'scores': None,
            'landmarks': None,
            'faces': None
        }
        print("人脸分析系统初始化完成。")
    @property
    def last_detection_results(self) -> dict:
        """
        获取最后一次detect_align的完整结果
        Returns:
            dict: 包含以下键的字典:
                - 'faces': 对齐后的人脸图像列表
                - 'boxes': 人脸框坐标列表
                - 'scores': 置信度列表
                - 'landmarks': 关键点列表
        """
        return self._last_detection

    def _validate_dataset(self, dataset_path: Path) -> bool:
        """
        验证人物数据集格式是否符合要求
        Args:
            dataset_path: 数据集路径
        Returns:
            bool: 是否验证通过
        """
        if not dataset_path.exists():
            raise ValueError(f"数据集路径不存在: {dataset_path}")

        if not dataset_path.is_dir():
            raise ValueError("数据集路径必须是一个目录")

        valid = True
        for person_dir in dataset_path.iterdir():
            if not person_dir.is_dir():
                print(f"警告: 发现非目录文件 {person_dir.name}，将被忽略")
                continue

            if not any(img.suffix.lower() in ['.jpg', '.jpeg', '.png'] for img in person_dir.iterdir()):
                print(f"错误: 人物目录 {person_dir.name} 中没有有效的图片文件")
                valid = False

        return valid

    def init_facebank_from_dataset(self, dataset_path: str):
        """
            从人物相片集初始化人脸库
            Args:
                dataset_path: 人物相片集路径
                conf: 配置对象(可选)
        """
        dataset_path = Path(dataset_path)
        if not self._validate_dataset(dataset_path):
            raise ValueError("数据集格式验证失败，请检查数据集结构")

        from facelib.InsightFace.models.utils import init_facebank, update_facebank
        conf.facebank_path = self.facebank_path
        # 初始化人脸库
        init_facebank(conf)

        # 遍历数据集目录
        dataset_path = Path(dataset_path)
        for person_dir in dataset_path.iterdir():
            if person_dir.is_dir():
                # 添加每个人的人脸数据
                add_from_folder(conf.facebank_path, str(person_dir), person_dir.name)

        # 更新人脸库
        update_facebank(conf, self.recognizer.model, self.detector)

    def detect(self, img: np.ndarray) -> Tuple[List, List, List, np.ndarray]:
        """基础人脸检测"""
        self._last_operations['detection'] = True
        boxes, scores, landmarks = self.detector.detect_faces(img)
        self._last_detection = {
            'boxes': boxes,
            'scores': scores,
            'landmarks': landmarks,
            'faces': None
        }
        return boxes, scores, landmarks, self.render(img)

    def align_face(self, img: np.ndarray) -> Tuple[List[np.ndarray], List, List, List, np.ndarray]:
        """
        人脸检测和对齐
        Args:
            img: 输入图像(OpenCV格式)
        Returns:
            faces: 对齐后的人脸图像列表
            boxes: 人脸框坐标列表
            scores: 置信度列表
            landmarks: 关键点列表
        """
        faces, boxes, scores, landmarks = self.detector.detect_align(img)
        self._last_detection = {
            'boxes': boxes,
            'scores': scores,
            'landmarks': landmarks,
            'faces': faces
        }
        return faces, boxes, scores, landmarks, self.render(img)

    def align_face_from_box(self, img: np.ndarray, box: List) -> Optional[np.ndarray]:
        """
        根据已知框对齐单个人脸
        Args:
            img: 原始图像
            box: 人脸框坐标[x1,y1,x2,y2]
        Returns:
            对齐后的人脸图像
        """
        # 获取参考关键点
        ref_pts = self.detector.ref_pts

        # 检测关键点
        _, _, landmarks = self.detector.detect_faces(img)
        if len(landmarks) == 0:
            return None

        # 计算变换矩阵
        trans = self.detector.trans
        trans.estimate(landmarks[0].cpu().numpy(), ref_pts)

        # 执行对齐
        aligned_face = cv2.warpAffine(img, trans.params[0:2, :], self.detector.out_size)
        return aligned_face

    def recognize(self, img: np.ndarray) -> Tuple[List[str], np.ndarray]:
        recognition_results = self._recognize(img)
        rendered_img = self.render(img, recognition_results)
        name_list = [result['identity'] for result in recognition_results]  # 提取名称列表
        return name_list, rendered_img

    def _recognize(self, img: np.ndarray) -> List[Dict]:
        """人脸识别"""
        self._last_operations['recognition'] = True
        faces, boxes, scores, landmarks = self.detector.detect_align(img)
        self._last_detection = {
            'boxes': boxes,
            'scores': scores,
            'landmarks': landmarks,
            'faces': faces
        }
        targets, names = self._load_facebank()
        results, _ = self.recognizer.infer(faces, targets)
        recognition_results = [{'identity': names[r + 1]} for r in results]
        return recognition_results

    def analyze_emotion(self, img: np.ndarray) -> Tuple[List[str], np.ndarray]:
        """情绪分析"""
        emotion_results = self._analyze_emotion(img)
        emotion_list = [result['emotion'] for result in emotion_results]  # 提取情绪字符串列表
        rendered_img = self.render(img, emotion_results)
        return emotion_list, rendered_img

    def _analyze_emotion(self, img: np.ndarray) -> List[Dict]:
        """情绪分析"""
        self._last_operations['emotion'] = True
        faces, boxes, scores, landmarks = self.emotion_detector.detect_align(img)
        self._last_detection = {
            'boxes': boxes,
            'scores': scores,
            'landmarks': landmarks,
            'faces': faces
        }
        emotions, _ = self.emotion.detect_emotion(faces)
        emotion_results = [{'emotion': e} for e in emotions]
        return emotion_results

    def analyze_age_gender(self, img: np.ndarray) -> Tuple[List[Tuple[str, int]], np.ndarray]:
        """年龄性别分析"""
        age_gender_results = self._analyze_age_gender(img)
        age_gender_list = [[result['gender'], result['age']] for result in age_gender_results]  # 转换为[gender, age]列表
        rendered_img = self.render(img, age_gender_results)
        return age_gender_list, rendered_img

    def _analyze_age_gender(self, img: np.ndarray) -> List[Dict]:
        """年龄性别分析内部实现"""
        self._last_operations['age_gender'] = True
        faces, boxes, scores, landmarks = self.detector.detect_align(img)
        self._last_detection = {
            'boxes': boxes,
            'scores': scores,
            'landmarks': landmarks,
            'faces': faces
        }
        genders, ages = self.age_gender.detect(faces)
        return [{'gender': g, 'age': a} for g, a in zip(genders, ages)]

    def full_analysis(self, img: np.ndarray) -> Tuple[List[Dict], np.ndarray]:
        """完整分析流程"""
        # 检测人脸
        faces, boxes, scores, landmarks = self.detector.detect_align(img)
        if len(faces) == 0:
            return [], img.copy()

        # 并行执行各分析任务
        recognition = self._recognize(img)
        emotion = self._analyze_emotion(img)
        age_gender = self._analyze_age_gender(img)

        # 合并结果
        results = []
        for i in range(len(faces)):
            result = {
                'box': boxes[i],
                'landmark': landmarks[i],
                **recognition[i],
                **emotion[i],
                **age_gender[i]
            }
            results.append(result)

        rendered_img = self.render(img, results)
        return results, rendered_img

    def render(self, img: np.ndarray, results: List[Dict] = None) -> np.ndarray:
        """智能渲染分析结果
        Args:
            img: 输入图像
            results: 分析结果列表，如果为None则只渲染检测框和关键点
        """
        rendered = img.copy()
        if self._last_detection['boxes'] is None or len(self._last_detection['boxes']) == 0:
            return rendered

        # 如果没有提供results，则只渲染检测框和关键点
        if results is None:
            for i in range(len(self._last_detection['boxes'])):
                box = [int(x) for x in self._last_detection['boxes'][i]]
                cv2.rectangle(rendered, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
                for pt in self._last_detection['landmarks'][i]:
                    cv2.circle(rendered, (int(pt[0]), int(pt[1])), 2, (0, 0, 255), -1)
            return rendered

        # 原有渲染逻辑
        for i, face in enumerate(results):
            box = [int(x) for x in self._last_detection['boxes'][i]]
            cv2.rectangle(rendered, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
            for pt in self._last_detection['landmarks'][i]:
                cv2.circle(rendered, (int(pt[0]), int(pt[1])), 2, (0, 0, 255), -1)

            texts = []
            if 'identity' in face:
                texts.append(f"Name: {face['identity']}")
            if 'emotion' in face:
                texts.append(f"Emo: {face['emotion']}")
            if 'age' in face and 'gender' in face:
                texts.append(f"{face['gender']} {face['age']}")

            for j, text in enumerate(texts):
                cv2.putText(rendered, text, (box[0], box[1] - 10 - j * 20),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)

        return rendered

    def _load_facebank(self) -> Tuple[torch.Tensor, List[str]]:
        """加载人脸特征库"""
        if not (self.facebank_path/'facebank.pth').exists():
            return torch.zeros(0), ['Unknown']
        embs = torch.load(self.facebank_path/'facebank.pth')
        names = np.load(self.facebank_path/'names.npy')
        return embs.to(self.device), names.tolist()
